package jtg.generator;

import jtg.exception.ConstraintSolveException;
import jtg.graphics.SootCFG; // 引用了用于获取控制流图的工具类。
import jtg.optimizer.Optimizer;
import jtg.solver.Z3Solver; // 引用了用于 Z3 约束求解的工具类。
import jtg.visualizer.Visualizer; // 引用了用于可视化的工具类。
import soot.Body;
import soot.Local;
import soot.Unit;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JIdentityStmt;
import soot.jimple.internal.JIfStmt;
import soot.jimple.internal.JReturnStmt;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

/**
 * SimpleGenerator：一个简单的测试用例生成器，通过使用 Soot 进行静态分析，结合 Z3Solver 进行约束求解，生成测试用例输入。
 * */
public class SimpleGenerator {
    private String clsPath; // 存储类路径
    private String clsName; // 类名
    private String mtdName; // 方法名
    private UnitGraph ug; // 存储方法的控制流图
    private Body body; // 存储方法体

    private Optimizer optimizer;

    private ArrayList<String> stepConditions;


    public SimpleGenerator(String className, String methodName) {
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new SimpleGenerator(defaultClsPath, className, methodName);
    }

    // 构造函数，获取方法体和控制流图
    public SimpleGenerator(String classPath, String className, String methodName) {
        clsPath = classPath;
        clsName = className;
        mtdName = methodName;
        ug = SootCFG.getMethodCFG(clsPath, clsName, mtdName);
        body = SootCFG.getMethodBody(clsPath, clsName, mtdName);
        optimizer = new Optimizer(body);
    }

    // 绘制控制流图
    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }


    // 获取 Jimple 增加的局部变量（以 "$" 开头），不是被测代码真正的变量。
    private List<Local> getJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<Local>();
        for (Local l : body.getLocals()) {
            if (l.toString().startsWith("$")) jimpleVars.add(l);
        }
        return jimpleVars;
    }

    private List<Local> getParameter() {
        ArrayList<Local> paras = new ArrayList<Local>();
        for (Local para : body.getParameterLocals()) {
            paras.add(para);
        }
        return paras;
    }

    // 生成测试用例输入
    public List<String> generate() {

        List<Unit> path = null;
        ArrayList<String> testSet = null;
        String pathConstraint = "";

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        try {
            testSet = new ArrayList<String>();
            // 遍历控制流图的头和尾，获取路径
            for (Unit h : ug.getHeads())
                for (Unit t : ug.getTails()) {
                    path = ug.getExtendedBasicBlockPathBetween(h, t);
                    System.out.println("The path is: " + path.toString());

                    /* 调用 calPathConstraint 方法计算路径约束：
                    在测试用例生成和程序分析的上下文中，路径约束通常用于描述在执行特定的程序路径时必须满足的条件。
                    这些条件可以包括分支语句的条件、循环语句的循环不变式、以及程序中的各种约束条件。 */
                    pathConstraint = calPathConstraint(path);

                    // 生成器会根据路径约束来构建测试用例，确保生成的输入满足路径上的条件。
                    // 如果路径约束为空字符串，表示路径约束为恒真。
                    if (pathConstraint.isEmpty()) testSet.add(randomTC(body.getParameterLocals()));
                    pathConstraint = optimizer.constraintOptimizer(stepConditions,path);
                    pathConstraint = pathConstraint.replace("$","");
                    System.out.println("The corresponding path constraint is: " + pathConstraint);
                    if (!pathConstraint.isEmpty())
                    {

                        String testItem = solve(pathConstraint);// solve方法 ==> Z3Solver根据路径约束生成测试用例
//                        testSet.add(testItem);
//                        System.out.println("begin");
//                        System.out.println(resultOptimize(testItem,body.getParameterLocals()));
//                        System.out.println("end");
                        testItem = optimizer.resultOptimize(testItem,path,pathConstraint);
                        testSet.add(testItem); //用resultOptimize优化器优化结果
//                        testSet.add(resultOptimize(testItem,body.getParameterLocals())); //用resultOptimize优化器优化结果
                    }
                }
        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
        }
        if (!testSet.isEmpty()) {
            System.out.println("");
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) " + tc);
                count++;
            }
        }
        return testSet;
    }

    // 根据给定的路径计算路径约束（只处理的整型）。
    public String calPathConstraint(List<Unit> path) {

        // jVars会获取执行路径下的变量（不是其真正的变量，可以理解为指令级别的临时变量）（局部变量是以$开头的）
        // 例如：[$i1, $i2, $i3, $r0]
        List<Local> jVars = getJVars();


        String pathConstraint = "";
        String expectedResult = ""; // 记录该条执行路径下的返回结果

        /*
        Jimple Vars（Jimple 变量）是指在 Jimple 语言中使用的变量。Jimple 是 Soot 工具包中的一种中间表示语言，用于表示 Java 字节码的简化版本。
        Jimple Vars 具有以下特点：
        1.局部性： Jimple Vars 是局部变量，它们的作用范围仅限于其声明所在的方法内部。
        2.虚拟寄存器： 在 Jimple 中，Jimple Vars 通常用虚拟寄存器来表示。
        3.语句级别： Jimple Vars 在 Jimple 语句中使用，例如，赋值语句、条件语句等。
        Jimple Vars 的名称以 $ 开头的字符串，例如 $i0、$r1 等。
         */
        HashMap<String, String> assignList = new HashMap<>(); // 用hash记录赋值语句的操作，key值为要赋值的变量，value是对key值（变量）的具体赋值操作
        ArrayList<String> stepConditionsWithJimpleVars = new ArrayList<String>(); // (可能带有Jimple变量)记录执行路径中满足条件判断的情况
        stepConditions = new ArrayList<String>();// (不带有Jimple变量)记录执行路径中满足条件判断的情况

        /*
        遍历执行路径
        该部分会对三种语句进行不同的处理：赋值语句、条件语句、return语句
         */
        for (Unit stmt : path) {
            /*
            =======================处理赋值语句=======================
            stmt ==> $i1 = i0 % 15

            判断变量 stmt 是否属于 JAssignStmt 类型或其子类型的实例。
            instanceof 是一个运算符，语法为：引用类型变量(object) instanceof 类(class)
            判断前面的对象是否属于后面的类，或者属于其子类；如果是，返回 true，不是返回 false；
            stmt 是 Unit 类型的对象，而 JAssignStmt 是 Unit 的一个子类。
            JAssignStmt 表示了一个 Jimple 赋值语句，其中一个值被分配给一个变量。
            ========================================================
             */
            if (stmt instanceof JAssignStmt) {
                // 如果当前语句是赋值语句，它会将左操作数和右操作数添加到 assignList 中，用于后续的路径约束计算。
                // 例如：$i1 = i0 % 15
                // ((JAssignStmt) stmt).getLeftOp().toString() ==> $i1
                // ((JAssignStmt) stmt).getRightOp().toString()==> i0 % 15
                assignList.put(((JAssignStmt) stmt).getLeftOp().toString(), ((JAssignStmt) stmt).getRightOp().toString());
                continue;
            }

            /*
            =======================处理条件语句=======================
            stmt ==> if $i1 != 0 goto $i2 = i0 % 5
            */
            if (stmt instanceof JIfStmt) {

                String ifstms = ((JIfStmt) stmt).getCondition().toString();// ifstms ==> $i1 != 0

                // 下面判断该条件判断语句在当前路径下是真是假？

                // step1: 获取条件判断语句下一条Unit
                int nextUnitIndex = path.indexOf(stmt) + 1;
                Unit nextUnit = path.get(nextUnitIndex);

                // step2: 判断当前条件是真是假？
                /*
                如果ifstmt的后继语句不是ifstmt中goto语句，说明ifstmt中的条件为假
                例如，下面的例子是条件为真的情况：
                    if $i1 != 0 goto $i2 = i0 % 5, $i2 = i0 % 5；
                    其中，((JIfStmt) stmt).getTarget() =$i2 = i0 % 5， nextUnit=$i2 = i0 % 5
                条件为假的情况：
                    if $i2 != 0 goto $i3 = i0 % 3, return "Buzz"

                ((JIfStmt) stmt).getTarget()==>获取条件判断语句中goto关键字之后的部分
                 */
                if (!((JIfStmt) stmt).getTarget().equals(nextUnit))
                    ifstms = "!( " + ifstms + " )";

                // 如果条件为真
                else
                    ifstms = "( " + ifstms + " )";
                stepConditionsWithJimpleVars.add(ifstms);
                continue;
            }

            /*
            =======================处理返回语句=======================
            expectedResult==>记录return的内容
             */
            if (stmt instanceof JReturnStmt) {
                expectedResult = stmt.toString().replace("return", "").trim();
            }
        }
        System.out.println("The step conditions with JimpleVars are: " + stepConditionsWithJimpleVars);

        // bug 没有考虑jVars为空的情况，jVars为空表示没有使用局部变量
        // TODO 没有考虑jVars为空的情况，即没有使用局部变量
        if (jVars.size() != 0) {
            for (String cond : stepConditionsWithJimpleVars) {
                //替换条件里的Jimple变量
                for (Local lv : jVars) {
                    if (cond.contains(lv.toString())) {
                        // 将路径中的临时变量替换为对应的赋值操作
                        // replace ==> 替换字符串中所有匹配的子字符串
                        stepConditions.add(cond.replace(lv.toString(), assignList.get(lv.toString()).trim()));
                        // 使用 trim()，可以确保右操作数的前后空白被移除，从而生成更干净、更可靠的路径约束。
                    }
                }
            }
        } else
            stepConditions = stepConditionsWithJimpleVars;

        if (stepConditions.isEmpty())
            return "";
        pathConstraint = stepConditions.get(0);
        int i = 1;
        while (i < stepConditions.size()) {
            pathConstraint = pathConstraint + " && " + stepConditions.get(i);
            i++;
        }
        System.out.println("The path expression is: " + pathConstraint);
        return pathConstraint;
    }

    // Z3Solver根据路径约束生成测试用例
    public String solve(String pathConstraint) throws Exception {
        return Z3Solver.solve(pathConstraint);
    }

    // 没有路径约束，生成random测试用例
    // input ==> 路径中要使用到的参数
    // output ==> 每个参数根据类型进行随机生成测试值
    public String randomTC(List<Local> parameters) {
        System.out.println("-----------------------------------------------------------------------------------");
        System.out.println(parameters);

        String varName;
        String varValue = "";
        String testinput = "";

        for (Local para : parameters) {
            varName = para.getName();
            if ("int".equals(para.getType().toString())) {
                varValue = String.valueOf((int)(Math.random() * 10));
            }
            if ("String".equals(para.getType().toString())) {
                varValue = "abc";
            }
            if ("char".equals(para.getType().toString())){
                varValue = String.valueOf('0');
            }
            //其它的基本类型没写
            // TODO
            testinput = testinput + " " + varName + "=" + varValue;
        }
        return testinput;
    }

}
